import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
import math

def process_image(img_path):
    try:
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        if img is None:
            print(f"无法读取图片: {img_path}")
            return None
        return img
    except Exception as e:
        print(f"处理图片时出错 {img_path}: {str(e)}")
        return None

def process_batch(args):
    image_paths, batch_id, output_folder, width, height = args
    output_video = os.path.join(output_folder, f"temp_video_{batch_id}.mp4")
    
    # 处理这一批次的图片
    images = []
    for img_path in tqdm(image_paths, desc=f"处理批次 {batch_id}", ncols=100):
        img = process_image(img_path)
        if img is not None:
            images.append(img)
    
    if not images:
        return None
        
    # 创建临时视频
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video, fourcc, 4000.0, (width, height))
    
    for img in images:
        out.write(img)
    
    out.release()
    return output_video

def merge_videos(video_files, output_video):
    # 读取第一个视频获取信息
    cap = cv2.VideoCapture(video_files[0])
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    cap.release()
    
    # 创建最终视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video, fourcc, 60.0, (width, height))
    
    # 依次读取每个临时视频并写入
    for video_file in tqdm(video_files, desc="合并视频", ncols=100):
        cap = cv2.VideoCapture(video_file)
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            out.write(frame)
        cap.release()
    
    out.release()

def main():
    # 设置输入输出路径
    input_folder = r"F:\边缘处理结果"
    output_folder = os.path.join(input_folder, "temp_videos")
    final_output = os.path.join(input_folder, "output.mp4")
    
    # 创建临时文件夹
    os.makedirs(output_folder, exist_ok=True)
    
    # 获取所有bmp文件
    try:
        image_files = [f for f in os.listdir(input_folder) if f.lower().endswith('.bmp')]
    except Exception as e:
        print(f"读取文件夹出错: {str(e)}")
        return
    
    if not image_files:
        print("文件夹中没有找到.bmp文件！")
        return
    
    image_files.sort()
    image_paths = [os.path.join(input_folder, f) for f in image_files]
    
    # 读取一张图片获取尺寸
    sample_img = process_image(image_paths[0])
    if sample_img is None:
        print("无法读取样本图片！")
        return
    height, width = sample_img.shape[:2]
    
    # 分批处理
    batch_size = 1000  # 每批处理1000张图片
    num_batches = math.ceil(len(image_paths) / batch_size)
    
    # 准备批次参数
    batch_args = []
    for i in range(num_batches):
        start_idx = i * batch_size
        end_idx = min((i + 1) * batch_size, len(image_paths))
        batch_args.append((
            image_paths[start_idx:end_idx],
            i,
            output_folder,
            width,
            height
        ))
    
    # 使用进程池处理批次
    num_cores = 4
    pool = Pool(processes=num_cores)
    
    print(f"使用 {num_cores} 个CPU核心进行处理...")
    temp_videos = list(pool.imap(process_batch, batch_args))
    
    pool.close()
    pool.join()
    
    # 过滤掉None值
    temp_videos = [v for v in temp_videos if v is not None]
    
    if not temp_videos:
        print("没有生成任何临时视频！")
        return
    
    # 合并视频
    print("正在合并视频...")
    merge_videos(temp_videos, final_output)
    
    # 清理临时文件
    print("清理临时文件...")
    for video_file in temp_videos:
        try:
            os.remove(video_file)
        except Exception as e:
            print(f"删除临时文件失败 {video_file}: {str(e)}")
    
    try:
        os.rmdir(output_folder)
    except Exception as e:
        print(f"删除临时文件夹失败: {str(e)}")
    
    print(f"视频已生成: {final_output}")
    total_frames = len(image_paths)
    print(f"视频时长: {total_frames/60:.2f} 秒")

if __name__ == "__main__":
    main()
